#!/usr/bin/python

"""
    Script to show LAG and LAG member status in a summary view
    Example of the output:
    acsadmin@sonic:~$ teamshow
    Flags: A - active, I - inactive, Up - up, Dw - down, N/A - Not Available,
           S - selected, D - deselected, * - not synced
     No.  Team Dev       Protocol    Ports
    -----  -------------  ----------  ---------------------------
        0  PortChannel0   LACP(A)(Up)     Ethernet0(D) Ethernet4(S)
        8  PortChannel8   LACP(A)(Up)     Ethernet8(S) Ethernet12(S)
       16  PortChannel16  LACP(A)(Up)     Ethernet20(S) Ethernet16(S)
       24  PortChannel24  LACP(A)(Dw)     Ethernet28(S) Ethernet24(S)
       32  PortChannel32  LACP(A)(Up)     Ethernet32(S) Ethernet36(S)
       40  PortChannel40  LACP(A)(Dw)     Ethernet44(S) Ethernet40(S)
       48  PortChannel48  LACP(A)(Up)     Ethernet52(S) Ethernet48(S)
       56  PortChannel56  LACP(A)(Dw)     Ethernet60(S) Ethernet56(S)

"""

import json
import os
import swsssdk
import subprocess
import sys
from tabulate import tabulate
from natsort import natsorted

PORT_CHANNEL_APPL_TABLE_PREFIX = "LAG_TABLE:"
PORT_CHANNEL_CFG_TABLE_PREFIX = "PORTCHANNEL|"
PORT_CHANNEL_STATUS_FIELD = "oper_status"

PORT_CHANNEL_MEMBER_APPL_TABLE_PREFIX = "LAG_MEMBER_TABLE:"
PORT_CHANNEL_MEMBER_STATUS_FIELD = "status"

class Teamshow(object):
    def __init__(self):
        self.teams = []
        self.teamsraw = {}
        self.summary = {}
        self.err = None
        # setup db connection
        self.db = swsssdk.SonicV2Connector(host="127.0.0.1")
        self.db.connect(self.db.APPL_DB)
        self.db.connect(self.db.CONFIG_DB)

    def get_portchannel_names(self):
        """
            Get the portchannel names from database.
        """
        team_keys = self.db.keys(self.db.CONFIG_DB, PORT_CHANNEL_CFG_TABLE_PREFIX+"*")
        if team_keys is None:
            return
        self.teams = [key[len(PORT_CHANNEL_CFG_TABLE_PREFIX):] for key in team_keys]

    def get_portchannel_status(self, port_channel_name):
        """
            Get port channel status from database.
        """
        full_table_id = PORT_CHANNEL_APPL_TABLE_PREFIX + port_channel_name
        return self.db.get(self.db.APPL_DB, full_table_id, PORT_CHANNEL_STATUS_FIELD)

    def get_portchannel_member_status(self, port_channel_name, port_name):
        full_table_id = PORT_CHANNEL_MEMBER_APPL_TABLE_PREFIX + port_channel_name + ":" + port_name
        return self.db.get(self.db.APPL_DB, full_table_id, PORT_CHANNEL_MEMBER_STATUS_FIELD)

    def get_team_id(self, team):
        """
            Skip the 'PortChannel' prefix and extract the team id.
        """
        return team[11:]

    def get_teamdctl(self):
        """
            Get teams raw data from teamdctl.
            Command: 'teamdctl <teamdevname> state dump'.
        """
        for team in self.teams:
            teamdctl_cmd = 'teamdctl ' + team + ' state dump'
            p = subprocess.Popen(teamdctl_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
            (output, err) = p.communicate()
            rc = p.wait()
	    if rc == 0:
                self.teamsraw[self.get_team_id(team)] = output
            else:
                self.err = err

    def get_teamshow_result(self):
        """
             Get teamshow results by parsing the output of teamdctl and combining port channel status.
        """
        for team in self.teams:
            info = {}
            team_id = self.get_team_id(team)
            if team_id not in self.teamsraw:
                info['protocol'] = 'N/A'
                self.summary[team_id] = info
                self.summary[team_id]['ports'] = ''
                continue
            json_info = json.loads(self.teamsraw[team_id])
            info['protocol'] = json_info['setup']['runner_name'].upper()
            info['protocol'] += '(A)' if json_info['runner']['active'] else '(I)'
            portchannel_status = self.get_portchannel_status(team)
            if portchannel_status is None:
                info['protocol'] += '(N/A)'
            elif portchannel_status.lower() == 'up':
                info['protocol'] += '(Up)'
            elif portchannel_status.lower() == 'down':
                info['protocol'] += '(Dw)'
            else:
                info['protocol'] += '(N/A)'

            info['ports'] = ""
            if 'ports' not in json_info:
                info['ports'] = 'N/A'
            else:
                for port in json_info['ports']:
                    status = self.get_portchannel_member_status(team, port)
                    selected = json_info["ports"][port]["runner"]["selected"]

                    info["ports"] += port + "("
                    info["ports"] += "S" if selected else "D"
                    if status is None or (status == "enabled" and not selected) or (status == "disabled" and selected):
                        info["ports"] += "*"
                    info["ports"] += ") "

            self.summary[team_id] = info

    def display_summary(self):
        """
            Display the portchannel (team) summary.
        """
        print("Flags: A - active, I - inactive, Up - up, Dw - Down, N/A - not available,\n"
              "       S - selected, D - deselected, * - not synced")

        header = ['No.', 'Team Dev', 'Protocol', 'Ports']
        output = []
        for team_id in natsorted(self.summary):
            output.append([team_id, 'PortChannel'+team_id, self.summary[team_id]['protocol'], self.summary[team_id]['ports']])
        print tabulate(output, header)

def main():
    if os.geteuid() != 0:
        exit("This utility must be run as root")

    try:
        team = Teamshow()
        team.get_portchannel_names()
        team.get_teamdctl()
        team.get_teamshow_result()
        team.display_summary()
    except Exception as e:
        sys.exit(e.message)

if __name__ == "__main__":
    main()
